package com.zeratul.common.pool.decorator;

import com.zeratul.common.core.context.DataContextHolde;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;

import java.util.Map;

/**
 * 线程装饰器
 */
public class ZtTaskDecorator implements TaskDecorator {

    /**
     * 线程池任务装饰器，这里的作用是让日志的MDC上下文可以传递到线程池的任务中
     */
    @Override
    public Runnable decorate(Runnable runnable) {
        Map<String, String> context = MDC.getCopyOfContextMap();
        Map<String, Object> reqContext = DataContextHolde.getTraceContextOnCopy();
        return () -> {
            MDC.setContextMap(context);
            DataContextHolde.setTraceContext(reqContext);
            try {
                runnable.run();
            } finally {
                MDC.clear();
                DataContextHolde.clear();
            }
        };
    }
}
